{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-21T13:31:25.022339600Z",
     "start_time": "2023-11-21T13:31:25.016155200Z"
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c552da9d-36f9-43d3-ae1f-c54f78d3ff2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from msprof_analyze.advisor.interface.interface import Interface\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from prettytable import PrettyTable, ALL\n",
    "from textwrap import fill"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57d17a21205c3c5e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# 集群调优分析\n",
    "## 1. 集群分析的数据准备\n",
    "首先我们当前支持PyTorch多卡大模型的集群分析，您需要输入集群分析的profiling_path路径，例如：  \n",
    "--{profiling_path}  \n",
    "    -- xxxx_ascend_pt  \n",
    "    -- xxxx_ascend_pt  \n",
    "    -- xxxx_ascend_pt  \n",
    "    ......  \n",
    "    -- xxxx_ascend_pt  \n",
    "里面每张卡的profiling文件都是ascend_pt结尾的文件。  \n",
    "\n",
    "## 2. 集群分析解决的问题  \n",
    "当前的功能主要有四项：  \n",
    "1）. 识别多卡间的计算慢卡（根据计算时间等推断）  \n",
    "2）. 识别多卡间的通信慢现象（根据通信链路的带宽判断）  \n",
    "3）. 对多卡间的计算算子进行统计展示（识别不同卡的算子差异）  \n",
    "4）. 展示集群流水并行图（根据时间轴展示多卡间的计算和通信时间）  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "36b7a24cc7ca5da2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-21T12:53:38.379699800Z",
     "start_time": "2023-11-21T12:53:38.363755900Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# EDIT THE PROFILING DATA PATH\n",
    "cluster_path = r\"YOUR PROFILING PATH\"\n",
    "interface = Interface(profiling_path=cluster_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf832ac2e0dfa30f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## 1) 识别慢卡"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "40aac93278dd6e34",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-21T12:53:41.815599700Z",
     "start_time": "2023-11-21T12:53:41.783393700Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "slow_rank_result = interface.get_result(\"cluster\", \"slow_rank\", template_key='overall')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0e943b2a-37a6-4db6-9e70-235d397f1d39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "            <th>step</th>\n",
       "            <th>rank_id</th>\n",
       "            <th>compute(us)</th>\n",
       "            <th>communication(us)</th>\n",
       "            <th>free(us)</th>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "        <tr>\n",
       "            <td>5</td>\n",
       "            <td>0</td>\n",
       "            <td>1725149.7719999943</td>\n",
       "            <td>520820.886</td>\n",
       "            <td>2897916.4370001294</td>\n",
       "        </tr>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "+------+---------+--------------------+-------------------+--------------------+\n",
       "| step | rank_id |    compute(us)     | communication(us) |      free(us)      |\n",
       "+------+---------+--------------------+-------------------+--------------------+\n",
       "|  5   |    0    | 1725149.7719999943 |     520820.886    | 2897916.4370001294 |\n",
       "+------+---------+--------------------+-------------------+--------------------+"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "slow_rank_data = slow_rank_result.get(\"慢卡分析\")\n",
    "if slow_rank_data:\n",
    "    slow_rank_table = PrettyTable(slow_rank_data.get(\"headers\"))\n",
    "    for row in slow_rank_data.get(\"data\"):\n",
    "        row = [fill(str(element), width=80) for element in row]\n",
    "        slow_rank_table.add_row(row)\n",
    "    slow_rank_table.hrules = ALL\n",
    "    display(slow_rank_table[:16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "57a9b1c6-4127-47a2-8699-3c983950bd84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "            <th>category</th>\n",
       "            <th>description</th>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "        <tr>\n",
       "            <td>慢卡分析</td>\n",
       "            <td>没有慢节点问题</td>\n",
       "        </tr>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "+----------+----------------+\n",
       "| category |  description   |\n",
       "+----------+----------------+\n",
       "| 慢卡分析 | 没有慢节点问题 |\n",
       "+----------+----------------+"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "problems = slow_rank_result.get(\"问题综述\")\n",
    "headers = problems.get('headers')[:2]\n",
    "if problems: # 如果存在相关问题则获取相关问题检测描述及建议\n",
    "    problem_table = PrettyTable(headers)\n",
    "    for row in problems.get(\"data\"):\n",
    "        row = [fill(str(element), width=100) for element in row]\n",
    "        problem_table.add_row(row[:2])\n",
    "    display(problem_table)\n",
    "else:\n",
    "    print(\"There is no suggestion related to slow rank analysis.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3511befaff513e8e",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## 2）识别通信链路慢"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2a1e617d2a117125",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "slow_link_result = interface.get_result(\"cluster\", \"slow_link\", template_key='overall')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c8bca314-a8da-4a5b-985a-c36f00154552",
   "metadata": {},
   "outputs": [],
   "source": [
    "slow_link_data = slow_link_result.get(\"slow_link_analysis\")\n",
    "if slow_link_data:\n",
    "    slow_link_table = PrettyTable(slow_link_data.get(\"headers\"))\n",
    "    for row in slow_link_data.get(\"data\"):\n",
    "        for i in range(len(row)):\n",
    "            row[i] = fill(str(row[i]), width=60)\n",
    "        slow_link_table.add_row(row)\n",
    "    slow_link_table.hrules = ALL\n",
    "    display(slow_link_table[:16])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "77d6efa1-48e3-409f-82c4-3e2b3d868898",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <thead>\n",
       "        <tr>\n",
       "            <th>category</th>\n",
       "            <th>description</th>\n",
       "        </tr>\n",
       "    </thead>\n",
       "    <tbody>\n",
       "        <tr>\n",
       "            <td>慢卡分析</td>\n",
       "            <td>没有慢节点问题</td>\n",
       "        </tr>\n",
       "        <tr>\n",
       "            <td>慢链路分析</td>\n",
       "            <td>RDMA bandwidth(GB/s)：      平均值是 24.356，      但最大值是 24.356GB/s ，     最小值是 24.356GB/s。     差距为<br>0.0GB/s。  SDMA bandwidth(GB/s)：      平均值是 18.181，      但最大值是 18.181GB/s ，     最小值是 18.181GB/s。<br>差距为 0.0GB/s。</td>\n",
       "        </tr>\n",
       "    </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "+------------+------------------------------------------------------------------------------------------------------------------+\n",
       "|  category  |                                                   description                                                    |\n",
       "+------------+------------------------------------------------------------------------------------------------------------------+\n",
       "|  慢卡分析  |                                                  没有慢节点问题                                                  |\n",
       "| 慢链路分析 | RDMA bandwidth(GB/s)：      平均值是 24.356，      但最大值是 24.356GB/s ，     最小值是 24.356GB/s。     差距为 |\n",
       "|            | 0.0GB/s。  SDMA bandwidth(GB/s)：      平均值是 18.181，      但最大值是 18.181GB/s ，     最小值是 18.181GB/s。 |\n",
       "|            |                                                 差距为 0.0GB/s。                                                 |\n",
       "+------------+------------------------------------------------------------------------------------------------------------------+"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "problems = slow_link_result.get(\"问题综述\")\n",
    "headers = problems.get('headers')[:2]\n",
    "if problems: # 如果存在相关问题则获取相关问题检测描述及建议\n",
    "    problem_table = PrettyTable(headers)\n",
    "    for row in problems.get(\"data\"):\n",
    "        row = [fill(str(element), width=100) for element in row]\n",
    "        problem_table.add_row(row[:2])\n",
    "    display(problem_table)\n",
    "else:\n",
    "    print(\"There is no suggestion related to slow link analysis.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce27a1d3-1354-45f7-88d8-dcb8e438b2b2",
   "metadata": {},
   "source": [
    "## 3) 分布式卡上的kernel算子统计展示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "466a0f30-042c-492a-bbf2-a5a85b649f95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from msprof_analyze.advisor.advisor_backend.interface import Interface\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e05774e9-c47e-400f-8421-b4b71bcdcbc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "interface = Interface(cluster_path)\n",
    "dataset = interface.get_data('cluster', 'kernel')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e95b6849-1738-4975-929f-734edff5d1c1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>rank id</th>\n",
       "      <th>Name</th>\n",
       "      <th>Input Shapes</th>\n",
       "      <th>Input Data Types</th>\n",
       "      <th>Output Shapes</th>\n",
       "      <th>Duration(us)_mean</th>\n",
       "      <th>Duration(us)_var</th>\n",
       "      <th>Duration(us)_max</th>\n",
       "      <th>Duration(us)_min</th>\n",
       "      <th>Duration(us)_count</th>\n",
       "      <th>Duration(us)_sum</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>ApplyAdamW1</td>\n",
       "      <td>\"2048;2048;2048;;;;;;;;2048\"</td>\n",
       "      <td>FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOA...</td>\n",
       "      <td>\"2048;2048;2048\"</td>\n",
       "      <td>4.060000</td>\n",
       "      <td>0.003200</td>\n",
       "      <td>4.100</td>\n",
       "      <td>4.020</td>\n",
       "      <td>2</td>\n",
       "      <td>8.120</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>ApplyAdamW1</td>\n",
       "      <td>\"273890816;273890816;273890816;;;;;;;;273890816\"</td>\n",
       "      <td>FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOA...</td>\n",
       "      <td>\"273890816;273890816;273890816\"</td>\n",
       "      <td>7369.147000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>7369.147</td>\n",
       "      <td>7369.147</td>\n",
       "      <td>1</td>\n",
       "      <td>7369.147</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>SwiGlu</td>\n",
       "      <td>\"20480,10240\"</td>\n",
       "      <td>DT_BF16</td>\n",
       "      <td>\"20480,5120\"</td>\n",
       "      <td>600.245391</td>\n",
       "      <td>1004.557501</td>\n",
       "      <td>680.394</td>\n",
       "      <td>540.771</td>\n",
       "      <td>64</td>\n",
       "      <td>38415.705</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>SwiGluGrad</td>\n",
       "      <td>\"20480,5120;20480,10240\"</td>\n",
       "      <td>DT_BF16;DT_BF16</td>\n",
       "      <td>\"20480,10240\"</td>\n",
       "      <td>859.961437</td>\n",
       "      <td>83.216257</td>\n",
       "      <td>885.318</td>\n",
       "      <td>849.877</td>\n",
       "      <td>32</td>\n",
       "      <td>27518.766</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>aclnnAbs_AbsAiCore_Abs</td>\n",
       "      <td>\"8192,1,2048\"</td>\n",
       "      <td>DT_BF16</td>\n",
       "      <td>\"8192,1,2048\"</td>\n",
       "      <td>26.174406</td>\n",
       "      <td>10.995655</td>\n",
       "      <td>31.521</td>\n",
       "      <td>19.660</td>\n",
       "      <td>32</td>\n",
       "      <td>837.581</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>619</th>\n",
       "      <td>0</td>\n",
       "      <td>aclnnSubs_SubAiCore_Sub</td>\n",
       "      <td>\"8165,1;\"</td>\n",
       "      <td>INT64;INT64</td>\n",
       "      <td>\"8165,1\"</td>\n",
       "      <td>5.600000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>5.600</td>\n",
       "      <td>5.600</td>\n",
       "      <td>1</td>\n",
       "      <td>5.600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>620</th>\n",
       "      <td>0</td>\n",
       "      <td>aclnnSubs_SubAiCore_Sub</td>\n",
       "      <td>\"8168,1;\"</td>\n",
       "      <td>INT64;INT64</td>\n",
       "      <td>\"8168,1\"</td>\n",
       "      <td>5.620000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>5.620</td>\n",
       "      <td>5.620</td>\n",
       "      <td>1</td>\n",
       "      <td>5.620</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>621</th>\n",
       "      <td>0</td>\n",
       "      <td>aclnnSubs_SubAiCore_Sub</td>\n",
       "      <td>\"8;\"</td>\n",
       "      <td>INT32;INT32</td>\n",
       "      <td>\"8\"</td>\n",
       "      <td>1.580000</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1.580</td>\n",
       "      <td>1.580</td>\n",
       "      <td>1</td>\n",
       "      <td>1.580</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>622</th>\n",
       "      <td>0</td>\n",
       "      <td>aclnnTopk_CastAiCore_Cast</td>\n",
       "      <td>\"4096,5\"</td>\n",
       "      <td>INT32</td>\n",
       "      <td>\"4096,5\"</td>\n",
       "      <td>7.250000</td>\n",
       "      <td>0.010426</td>\n",
       "      <td>7.600</td>\n",
       "      <td>7.020</td>\n",
       "      <td>32</td>\n",
       "      <td>232.000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>623</th>\n",
       "      <td>0</td>\n",
       "      <td>aclnnTopk_TopkV2_TopKV2</td>\n",
       "      <td>\"4096,32;\"</td>\n",
       "      <td>FLOAT;INT32</td>\n",
       "      <td>\"4096,5;4096,5\"</td>\n",
       "      <td>39.940375</td>\n",
       "      <td>12.724348</td>\n",
       "      <td>45.021</td>\n",
       "      <td>35.941</td>\n",
       "      <td>32</td>\n",
       "      <td>1278.092</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>624 rows × 11 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     rank id                       Name  \\\n",
       "0          0                ApplyAdamW1   \n",
       "1          0                ApplyAdamW1   \n",
       "2          0                     SwiGlu   \n",
       "3          0                 SwiGluGrad   \n",
       "4          0     aclnnAbs_AbsAiCore_Abs   \n",
       "..       ...                        ...   \n",
       "619        0    aclnnSubs_SubAiCore_Sub   \n",
       "620        0    aclnnSubs_SubAiCore_Sub   \n",
       "621        0    aclnnSubs_SubAiCore_Sub   \n",
       "622        0  aclnnTopk_CastAiCore_Cast   \n",
       "623        0    aclnnTopk_TopkV2_TopKV2   \n",
       "\n",
       "                                         Input Shapes  \\\n",
       "0                        \"2048;2048;2048;;;;;;;;2048\"   \n",
       "1    \"273890816;273890816;273890816;;;;;;;;273890816\"   \n",
       "2                                       \"20480,10240\"   \n",
       "3                            \"20480,5120;20480,10240\"   \n",
       "4                                       \"8192,1,2048\"   \n",
       "..                                                ...   \n",
       "619                                         \"8165,1;\"   \n",
       "620                                         \"8168,1;\"   \n",
       "621                                              \"8;\"   \n",
       "622                                          \"4096,5\"   \n",
       "623                                        \"4096,32;\"   \n",
       "\n",
       "                                      Input Data Types  \\\n",
       "0    FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOA...   \n",
       "1    FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOAT;FLOA...   \n",
       "2                                              DT_BF16   \n",
       "3                                      DT_BF16;DT_BF16   \n",
       "4                                              DT_BF16   \n",
       "..                                                 ...   \n",
       "619                                        INT64;INT64   \n",
       "620                                        INT64;INT64   \n",
       "621                                        INT32;INT32   \n",
       "622                                              INT32   \n",
       "623                                        FLOAT;INT32   \n",
       "\n",
       "                       Output Shapes  Duration(us)_mean  Duration(us)_var  \\\n",
       "0                   \"2048;2048;2048\"           4.060000          0.003200   \n",
       "1    \"273890816;273890816;273890816\"        7369.147000               NaN   \n",
       "2                       \"20480,5120\"         600.245391       1004.557501   \n",
       "3                      \"20480,10240\"         859.961437         83.216257   \n",
       "4                      \"8192,1,2048\"          26.174406         10.995655   \n",
       "..                               ...                ...               ...   \n",
       "619                         \"8165,1\"           5.600000               NaN   \n",
       "620                         \"8168,1\"           5.620000               NaN   \n",
       "621                              \"8\"           1.580000               NaN   \n",
       "622                         \"4096,5\"           7.250000          0.010426   \n",
       "623                  \"4096,5;4096,5\"          39.940375         12.724348   \n",
       "\n",
       "     Duration(us)_max  Duration(us)_min  Duration(us)_count  Duration(us)_sum  \n",
       "0               4.100             4.020                   2             8.120  \n",
       "1            7369.147          7369.147                   1          7369.147  \n",
       "2             680.394           540.771                  64         38415.705  \n",
       "3             885.318           849.877                  32         27518.766  \n",
       "4              31.521            19.660                  32           837.581  \n",
       "..                ...               ...                 ...               ...  \n",
       "619             5.600             5.600                   1             5.600  \n",
       "620             5.620             5.620                   1             5.620  \n",
       "621             1.580             1.580                   1             1.580  \n",
       "622             7.600             7.020                  32           232.000  \n",
       "623            45.021            35.941                  32          1278.092  \n",
       "\n",
       "[624 rows x 11 columns]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "27b75df4-792b-43dc-aa5c-d3c265642c1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存到csv查看， 可修改保存路径\n",
    "dataset.to_csv('cluster_kernel_details.csv', index=False, sep='\\t')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae45826394463cc4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## 4) 展示集群流水并行图\n",
    "使用说明：  \n",
    "1）. 需要使用Ascend Torch Profiler采集数据，如果需要展示FP和BP需要将activities设置为采集CPU和NPU  \n",
    "2）. rank_ids为要展示的rank id列表，必选参数, 可视化顺序与rank_ids的顺序一致  \n",
    "3）. worker_num为多进程数量，可选参数，请根据机器配置调整，默认值为机器可用核心数的一半  \n",
    "4）. 如果没有采集CPU数据，则展示Stage和Bubble的流水图  \n",
    "5）. 生成的json文件可以在chrome trace中查看  \n",
    "\n",
    "示例图：\n",
    "![pipeline_view](../../profiler/msprof_analyze/test/resource/pipeline_view.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f34ecf5-5c4a-4bc0-a761-e6338e534bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# rank_ids为要呈现的rank id列表，必选参数\n",
    "# 可以使用列表推导式生成需要的rank_ids，最终展示顺序和rank_ids的顺序一致\n",
    "# worker_num为多进程数量，可选参数，请根据机器配置调整，默认值为机器可用核心数的一半\n",
    "dataset = interface.get_data(\"cluster\", \"pipeline\", rank_ids=[0, 1, 2, 3, 4, 5, 6, 7], worker_num=8)\n",
    "\n",
    "# 保存json数据，在chrome trace中查看\n",
    "with open(\"./pipeline_view.json\", \"w\") as f:\n",
    "    json.dump(dataset.get(\"data\", []), f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
